import cv2
import pygame
import sys
import time
import mediapipe as mp
# --- 配置和资源路径加载（支持PyInstaller打包） ---
import os

# --- MediaPipe Tasks API 别名 ---
BaseOptions = mp.tasks.BaseOptions
PoseLandmarker = mp.tasks.vision.PoseLandmarker
PoseLandmarkerOptions = mp.tasks.vision.PoseLandmarkerOptions
RunningMode = mp.tasks.vision.RunningMode

# 绘制实心爱心函数
def draw_heart(surface, x, y, size, color):
    # 两个半圆和下方三角形组合
    top_radius = size // 2
    # 左半圆
    pygame.draw.circle(surface, color, (x - top_radius, y - top_radius), top_radius)
    # 右半圆
    pygame.draw.circle(surface, color, (x + top_radius, y - top_radius), top_radius)
    # 三角形
    points = [(x - size, y - top_radius), (x + size, y - top_radius), (x, y + size)]
    pygame.draw.polygon(surface, color, points)

# 资源基本路径：PyInstaller打包时资源在 _MEIPASS，否则在脚本目录
if getattr(sys, 'frozen', False):
    base_path = sys._MEIPASS
else:
    base_path = os.path.dirname(__file__)

# --- 配置加载 ---
import json, os
# 配置文件存放在当前工作目录下，支持exe打包后写入
config_path = os.path.join(os.getcwd(), "config.json")
# 默认配置
config = {
    "CAM_INDEX": 0,
    "WINDOW_WIDTH": 1280,
    "WINDOW_HEIGHT": 720,
    "AVATAR_SIZE": 64,
    "HEART_RADIUS": 20,
    "RAISE_HEIGHT_THRESHOLD": 0.1,
    "INITIAL_HP": 10,
    "FIRE_DIST_THRESHOLD": 0.05,
    "DETECT_INTERVAL": 2
}
# 读取或生成配置文件
if os.path.exists(config_path):
    try:
        with open(config_path, 'r', encoding='utf-8') as f:
            user_cfg = json.load(f)
            config.update(user_cfg)
    except Exception as e:
        print(f"Warning: 无法读取 config.json，使用默认配置：{e}")
else:
    with open(config_path, 'w', encoding='utf-8') as f:
        json.dump(config, f, indent=4, ensure_ascii=False)

# 赋值配置
CAM_INDEX = config["CAM_INDEX"]
WINDOW_WIDTH = config["WINDOW_WIDTH"]
WINDOW_HEIGHT = config["WINDOW_HEIGHT"]
AVATAR_SIZE = config["AVATAR_SIZE"]
HEART_RADIUS = config["HEART_RADIUS"]
RAISE_HEIGHT_THRESHOLD = config["RAISE_HEIGHT_THRESHOLD"]
INITIAL_HP = config["INITIAL_HP"]
FIRE_DIST_THRESHOLD = config["FIRE_DIST_THRESHOLD"]
DETECT_INTERVAL = config.get("DETECT_INTERVAL", 1)  # 检测间隔帧数，1为每帧检测
# 初始化检测帧计数器
frame_count = 0
# 上次检测结果缓存
last_poses = []
RAISE_HEIGHT_THRESHOLD = 0.1  # 双手抬起高度阈值，Y轴归一化坐标
INITIAL_HP = 10  # 初始血量

# --- 游戏状态 ---
class GameState:
    WAIT_PLAYERS = 1
    WAIT_READY = 2
    IN_GAME = 3
    GAME_OVER = 4

# --- 初始化 ---
pygame.init()
screen = pygame.display.set_mode((WINDOW_WIDTH, WINDOW_HEIGHT))
pygame.display.set_caption("双人动作检测游戏")
font = pygame.font.SysFont("SimHei", 36)
clock = pygame.time.Clock()
cap = cv2.VideoCapture(CAM_INDEX)
if not cap.isOpened():
    print("无法打开摄像头")
    sys.exit(1)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, WINDOW_WIDTH)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, WINDOW_HEIGHT)

# --- MediaPipe PoseLandmarker 初始化 ---
# 模型文件路径（支持打包后资源路径）
model_path = os.path.join(base_path, "pose_landmarker_full.task")
base_options = BaseOptions(model_asset_path=model_path)
options = PoseLandmarkerOptions(
    base_options=base_options,
    running_mode=RunningMode.VIDEO,
    num_poses=2
)
pose_landmarker = PoseLandmarker.create_from_options(options)

# 当前状态
hp = [INITIAL_HP, INITIAL_HP]
from ball import EnergyBall
import numpy as np
balls = []
fire_flags = [False, False]
FIRE_DIST_THRESHOLD = 0.05  # 发射动作两手距离归一化阈值
state = GameState.WAIT_PLAYERS
avatars = [None, None]
ready_flags = [False, False]

# 主循环
while True:
    ret, frame = cap.read()
    if not ret:
        print("摄像头读取失败！")
        break

    # BGR -> RGB 并 resize
    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    frame = cv2.resize(frame, (WINDOW_WIDTH, WINDOW_HEIGHT))

    # 转换为 MediaPipe Image
    frame_count += 1
    do_detect = (frame_count == 1) or (frame_count % DETECT_INTERVAL == 0)

    mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=frame)

    # 多人姿态检测（可按 DETECT_INTERVAL 限制频率）
    timestamp_ms = int(time.time() * 1000)
    if do_detect:
        result = pose_landmarker.detect_for_video(mp_image, timestamp_ms)
        last_poses = result.pose_landmarks  # 注意：新版 API 是 pose_landmarks
    poses = last_poses or []

    # 背景绘制
    bg_surf = pygame.surfarray.make_surface(frame.swapaxes(0, 1))
    screen.blit(bg_surf, (0, 0))

    # 人数判断
    if len(poses) < 2:
        # 进入等待玩家状态时重置游戏参数
        if state != GameState.WAIT_PLAYERS:
            hp[:] = [INITIAL_HP, INITIAL_HP]
            balls.clear()
            fire_flags[:] = [False, False]
            ready_flags[:] = [False, False]
        state = GameState.WAIT_PLAYERS
    elif state == GameState.WAIT_PLAYERS:
        state = GameState.WAIT_READY

    # 排序：按每组pose平均x坐标
    if len(poses) >= 2:
        xs = [sum([lm.x for lm in p]) / len(p) for p in poses]
        indices = sorted(range(len(xs)), key=lambda i: xs[i])
        player_poses = [poses[indices[0]], poses[indices[1]]]
    else:
        player_poses = []

    # 头像截取与准备检测
    for pid, p in enumerate(player_poses):
        nose = p[0]  # landmark 0 = 鼻子
        cx, cy = int(nose.x * WINDOW_WIDTH), int(nose.y * WINDOW_HEIGHT)
        size = AVATAR_SIZE
        x0 = max(cx - size // 2, 0)
        y0 = max(cy - size // 2, 0)
        avatar_img = frame[y0:y0+size, x0:x0+size]
        if avatar_img.size > 0:
            surf = pygame.surfarray.make_surface(avatar_img.swapaxes(0, 1))
            avatars[pid] = pygame.transform.scale(surf, (size, size))

                # 准备检测：手腕与肩膀坐标
                # 准备检测：手腕与肩膀坐标
                # 准备检测：双手高举视为准备完成
        lwr, rwr = p[15], p[16]
        lsh, rsh = p[11], p[12]
        # 双手高举：手腕中点高于肩膀中点一个阈值
        shoulder_mid_y = (lsh.y + rsh.y) / 2
        wrist_mid_y = (lwr.y + rwr.y) / 2
        raised = wrist_mid_y < (shoulder_mid_y - RAISE_HEIGHT_THRESHOLD)
        ready_flags[pid] = raised

    # 绘制头像、胸口爱心攻击区和准备状态
    # 绘制胸口爱心攻击区
    for pid, p in enumerate(player_poses):
        lsh, rsh = p[11], p[12]
        cx = int(((lsh.x + rsh.x)/2) * WINDOW_WIDTH)
        cy = int(((lsh.y + rsh.y)/2) * WINDOW_HEIGHT)
        draw_heart(screen, cx, cy, HEART_RADIUS, (255,0,0))

    if len(avatars) >= 2 and avatars[0] and avatars[1]:
        screen.blit(avatars[0], (10, 10))
        screen.blit(avatars[1], (WINDOW_WIDTH-10-AVATAR_SIZE, 10))
    if state == GameState.WAIT_READY:
        # 等待准备提示：双手举过头顶
        STATUS = "请双手举过头顶以完成准备"
        # 重置小球和发射标记
        balls.clear()
        fire_flags[:] = [False, False]
        if all(ready_flags):
            state = GameState.IN_GAME
    elif state == GameState.WAIT_PLAYERS:
        STATUS = "请两名玩家进入画面"
    elif state == GameState.IN_GAME:
        STATUS = "游戏中(合拢双手发射气功波)"
        HEART_RADIUS = 20  # 碰撞判定半径
        # 发射检测与能量球管理
        for pid, p in enumerate(player_poses):
            lwr, rwr = p[15], p[16]
            dist_norm = ((lwr.x - rwr.x)**2 + (lwr.y - rwr.y)**2)**0.5
            if dist_norm < FIRE_DIST_THRESHOLD and not fire_flags[pid]:
                mx = ((lwr.x + rwr.x)/2) * WINDOW_WIDTH
                my = ((lwr.y + rwr.y)/2) * WINDOW_HEIGHT
                # 目标胸口方向
                lsh, rsh = p[11], p[12]
                cx = ((lsh.x + rsh.x)/2) * WINDOW_WIDTH
                cy = ((lsh.y + rsh.y)/2) * WINDOW_HEIGHT
                direction = (mx - cx, my - cy)
                balls.append(EnergyBall(pid+1, (mx, my), direction))
                fire_flags[pid] = True
            elif dist_norm >= FIRE_DIST_THRESHOLD:
                fire_flags[pid] = False
        # 更新和绘制小球
        for ball in list(balls):
            ball.update()
            ball.draw(screen)
            if ball.is_offscreen(WINDOW_WIDTH, WINDOW_HEIGHT):
                balls.remove(ball)
                continue
            # 碰撞检测
            targ = 1 if ball.owner_id == 1 else 0
            if len(player_poses) > targ:
                tp = player_poses[targ]
                lsh_t, rsh_t = tp[11], tp[12]
                cx_t = ((lsh_t.x + rsh_t.x)/2) * WINDOW_WIDTH
                cy_t = ((lsh_t.y + rsh_t.y)/2) * WINDOW_HEIGHT
                if ((ball.position[0]-cx_t)**2 + (ball.position[1]-cy_t)**2)**0.5 < ball.radius + HEART_RADIUS:
                    hp[targ] -= 1
                    balls.remove(ball)
                    if hp[targ] <= 0:
                        state = GameState.GAME_OVER
    else:
        STATUS = "状态: %s" % state

    # 绘制血量条和状态文字
    # 血量条尺寸
    BAR_WIDTH, BAR_HEIGHT = 200, 20
    # 玩家1血条
    pygame.draw.rect(screen, (128,128,128), (10, 80, BAR_WIDTH, BAR_HEIGHT))
    pygame.draw.rect(screen, (0,255,0), (10, 80, BAR_WIDTH * (hp[0]/INITIAL_HP), BAR_HEIGHT))
    # 玩家2血条
    pygame.draw.rect(screen, (128,128,128), (WINDOW_WIDTH-10-BAR_WIDTH, 80, BAR_WIDTH, BAR_HEIGHT))
    pygame.draw.rect(screen, (255,0,0), (WINDOW_WIDTH-10-BAR_WIDTH, 80, BAR_WIDTH * (hp[1]/INITIAL_HP), BAR_HEIGHT))
    # 血量数字
    hp_text1 = font.render(f"{hp[0]}", True, (255,255,0))
    screen.blit(hp_text1, (10 + BAR_WIDTH + 5, 80))
    hp_text2 = font.render(f"{hp[1]}", True, (255,255,0))
    screen.blit(hp_text2, (WINDOW_WIDTH-10-BAR_WIDTH-30, 80))

    # 状态文字
    text = font.render(STATUS, True, (255, 255, 0))
    tr = text.get_rect(center=(WINDOW_WIDTH // 2, WINDOW_HEIGHT - 30))
    screen.blit(text, tr)

    # GAME_OVER 界面
    if state == GameState.GAME_OVER:
        # 显示胜利信息及重启提示
        loser = 0 if hp[0] <= 0 else 1
        winner = 2 - loser  # 1 or 2
        end_text = font.render(f"玩家{winner} 获胜！按 R 重启", True, (255,255,0))
        et_rect = end_text.get_rect(center=(WINDOW_WIDTH//2, WINDOW_HEIGHT//2))
        screen.blit(end_text, et_rect)
    # 事件处理
    for event in pygame.event.get():
        if event.type == pygame.KEYDOWN and state == GameState.GAME_OVER and event.key == pygame.K_r:
            # 重置游戏
            hp = [INITIAL_HP, INITIAL_HP]
            balls.clear()
            fire_flags[:] = [False, False]
            ready_flags[:] = [False, False]
            state = GameState.WAIT_PLAYERS
            continue
        if event.type == pygame.QUIT:
            cap.release()
            pygame.quit()
            sys.exit(0)

    # 显示FPS
    fps = clock.get_fps()
    fps_surf = font.render(f"FPS: {fps:.1f}", True, (255,255,0))
    fps_rect = fps_surf.get_rect(topright=(WINDOW_WIDTH-10, 10))
    screen.blit(fps_surf, fps_rect)
    pygame.display.flip()
    clock.tick(30)
